# -*- coding: utf-8 -*-
"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/10/30 14:16
@Description: 
@Change     : 
@File       : DownLoader.py
===============================
"""
import sys
import platform


if platform.system().lower() == 'linux':
    sys.path.append("/home/zuo/26_DocumentDelivery/")
from elasticsearch import Elasticsearch, helpers
import os
import platform
import time
from datetime import datetime, timedelta
from loguru import logger
from difflib import SequenceMatcher
from concurrent.futures import ThreadPoolExecutor
from queue import Queue
from Config.Emails import emails
from Api.BatchAttachmentDownloader import DownloadEmailAttachment
from Api.mongo_client import mongo_client2
from Api.BaseApi import find_files_in_directory

es = Elasticsearch(["http://140.143.51.14:9200/"], http_auth=("robot", "ko0cXhepXD3FOLtq"))

Q = Queue()


def find_task_id_by_similarity(documents, title):
    for doc in documents:
        similarity = SequenceMatcher(None, doc['title'], title).ratio()
        if similarity > 0.8:
            return doc['task_id'], doc['title'], doc["doi"]
    return None, None, None


class DownLoader:
    def __init__(self, days=3):
        self.query = {'status': 0}
        self.projection = {'_id': 0, 'title': 1, "task_id": 1, "doi": 1}
        self.TASK_TABLE = "DocumentDeliveryTask"
        self.days = days

    def handle_task(self):
        while not Q.empty():
            try:
                info = Q.get()
                today = datetime.now().strftime('%Y-%m-%d')
                email_address = info['email_address']
                SAVE_PATH = fr'Email-Attachments\{email_address}\{today}' if platform.system().lower() != 'linux' else f'/mnt/DocumentDelivery/Applet/Email-Attachments/{email_address}/{today}'
                directory = [os.path.basename(_) for _ in find_files_in_directory(SAVE_PATH)]
                os.makedirs(SAVE_PATH, exist_ok=True)
                Date_Begin = (datetime.now() - timedelta(days=self.days)).strftime('%Y-%m-%d %H:%M')
                query = {'status': 2, 'email': email_address}
                document = mongo_client2.find_document(self.TASK_TABLE, query, projection=self.projection, Type=2)
                document = [{"task_id": _['task_id'], "title": _['title'], "doi": _['doi']} for _ in document[1] if document[0]]
                try:
                    downloader = DownloadEmailAttachment.DownloadEmailAttachment(email_server=info['server_address'], email_address=email_address, email_password=info['email_password'], save_dir=SAVE_PATH)
                    downloader.date_begin = Date_Begin
                    attachments, error_file_dist = downloader.download_attachments(directory)
                    # print(f"{info['email_address']}: {attachments}, {error_file_dist}")
                except Exception as e:
                    logger.error(f"{info['email_address']}下载失败：{e}")
                    continue
                else:
                    for subject, attachment in attachments.items():
                        task_id, title, doi = find_task_id_by_similarity(document, subject)
                        if not task_id:
                            continue
                        query = {'task_id': task_id}
                        local_file_path = attachment.replace("/mnt/DocumentDelivery", "") if platform.system().lower() == 'linux' else attachment.replace("Email-Attachments", "")
                        update = {'$set': {'status': 1, 'local_file_path': local_file_path, 'download_time': datetime.now()}}
                        update_bool, update_info = mongo_client2.update_document(col=self.TASK_TABLE, query=query, update=update)
                        if update_bool:
                            data = [{
                                "_index": 'oss-pdf2',
                                "_id": task_id,
                                "_source": {
                                    "title": title,
                                    "doi": doi,
                                    "pmid": "",
                                    "pmc": "",
                                    "path": local_file_path,
                                    "status1": 0,
                                    "status2": 0,
                                    "text_code": "zhDoc",
                                    "source": "AutoDownLoad",  # 自动化下载
                                    "text": "doi",
                                    "create_time": int(time.time()),
                                    "link": ""
                                }
                            }]
                            helpers.bulk(es, data, request_timeout=300)
                            logger.debug(f"{task_id} 文献传递成功")
                    for subject, attachment in error_file_dist.items():
                        task_id, title, doi = find_task_id_by_similarity(document, subject)
                        if not task_id:
                            continue
                        query = {'task_id': task_id}
                        mongo_client2.update_document(col=self.TASK_TABLE, query=query, update={'$set': {'status': -2, 'error_msg': f"附件格式检验未通过"}})
                        logger.error(f"{task_id} 文献附件格式检验未通过")
            except Exception as e:
                logger.error(f"下载失败：{e}")
                continue

    @staticmethod
    def to_queue():
        for info in emails:
            # if info['email_address'] == 'fm342897@163.com':
            Q.put(info)
        logger.info(f"任务队列初始化完成，共{Q.qsize()}个任务")

    def run(self, max_workers=10):
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            for i in range(max_workers):
                executor.submit(self.handle_task)


downloader = DownLoader()
if __name__ == '__main__':
    while True:
        try:
            downloader.to_queue()
            downloader.run(max_workers=10)
            # downloader.handle_task()
            time.sleep(60)
        except Exception as e:
            logger.error(f"程序出错：{e}")
            time.sleep(60)